import numpy as np
import torch.nn as nn

dout = 0.3  # dropout


class imgSiamese(nn.Module):
    """
        imgSiamese的模型实现
    """

    def __init__(self):
        super(imgSiamese, self).__init__()
        #   输入维度应该为batch*inchannel*N*N
        self.conv1 = nn.Conv2d(
            in_channels=1,  # 输入的feature map
            out_channels=64,  # 输出的feature map
            kernel_size=10,  # 卷积核尺寸
            stride=1,  # 卷积核步长
            padding=0,  # 不进行填充
        )
        self.conv2 = nn.Conv2d(64, 128, 7, 1, 0)
        self.conv3 = nn.Conv2d(128, 128, 4, 1, 0)
        self.conv4 = nn.Conv2d(128, 256, 4, 1, 0)
        self.dout1 = nn.Dropout2d(p=dout)
        self.dout2 = nn.Dropout2d(p=dout)
        self.dout3 = nn.Dropout2d(p=dout)
        self.dout4 = nn.Dropout2d(p=dout)
        self.dout5 = nn.Dropout(p=dout)
        self.dout6 = nn.Dropout(p=dout)
        self.fc1 = nn.Linear(
            in_features=256*6*6,  # 输入特征
            out_features=4096,  # 输出特证数
        )
        self.fc2 = nn.Linear(
            in_features=4096,  # 输入特征
            out_features=2048,  # 输出特证数
        )
        self.fc3 = nn.Linear(
            in_features=2048,  # 输入特征
            out_features=1024,  # 输出特证数
        )
        # self.fc4 = nn.Linear(
        #     in_features=1,
        #     out_features=1,
        # )
        nn.init.xavier_uniform_(self.conv1.weight)
        nn.init.xavier_uniform_(self.conv2.weight)
        nn.init.xavier_uniform_(self.conv3.weight)
        nn.init.xavier_uniform_(self.conv4.weight)
        nn.init.xavier_uniform_(self.fc1.weight)
        nn.init.xavier_uniform_(self.fc2.weight)
        self.sig = nn.Sigmoid()
        self.rl = nn.ReLU()
        self.mpool = nn.MaxPool2d(
            kernel_size=2,  # 平均值池化层,使用 2*2
            stride=2,  # 池化步长为2
        )
        self.cos = nn.CosineSimilarity(dim=1, eps=1e-6)
        self.dropout = nn.Dropout(p=dout)

    def get_representation(self, x):
        x = self.conv1(x)
        x = self.dout1(x)
        x = self.rl(x)
        x = self.mpool(x)

        x = self.conv2(x)
        x = self.dout2(x)
        x = self.rl(x)
        x = self.mpool(x)

        x = self.conv3(x)
        x = self.dout3(x)
        x = self.rl(x)
        x = self.mpool(x)

        x = self.conv4(x)
        x = self.dout4(x)
        x = x.view(x.size(0), -1)
        x = self.fc1(x)
        x = self.dout5(x)
        x = self.fc2(x)
        x = self.dout6(x)
        x = self.fc3(x)
        return x

    def forward(self, x, y):
        x = self.get_representation(x)
        y = self.get_representation(y)
        res = self.cos(x, y)
        return res
